import gym
from gym import spaces
import numpy as np

class BlackjackEnv(gym.Env):
    """
    A simple implementation of the Blackjack game as an OpenAI Gym environment.
    """

    def __init__(self):
        self.action_space = spaces.Discrete(2)
        self.observation_space = spaces.Tuple((
            spaces.Discrete(32),
            spaces.Discrete(11),
            spaces.Discrete(2)
        ))
        self.reward_range = (-1.0, 1.0)
        self.reset()

    def reset(self):
        self.dealer_sum = np.random.randint(12, 22)
        self.player_sum = np.random.randint(12, 22)
        self.usable_ace = False
        return (self.player_sum, self.dealer_sum, self.usable_ace)

    def step(self, action):
        if action == 0:
            # Stick
            done = True
            while self.dealer_sum < 17:
                self.dealer_sum += self.get_card()
            reward = self.get_reward()
        else:
            # Hit
            self.player_sum += self.get_card()
            done = self.check_bust()
            reward = self.get_reward(done)
        return (self.player_sum, self.dealer_sum, self.usable_ace), reward, done, {}

    def get_card(self):
        card = np.random.randint(1, 14)
        if card > 10:
            card = 10
        elif card == 1:
            if self.player_sum + 11 <= 21:
                card = 11
                self.usable_ace = True
            else:
                card = 1
        return card

    def check_bust(self):
        if self.player_sum > 21:
            if self.usable_ace:
                self.player_sum -= 10
                self.usable_ace = False
            else:
                return True
        return False

    def get_reward(self, done=False):
        if done:
            if self.player_sum > 21:
                return -1.0
            elif self.dealer_sum > 21:
                return 1.0
            elif self.player_sum > self.dealer_sum:
                return 1.0
            elif self.player_sum < self.dealer_sum:
                return -1.0
            else:
                return 0.0
        else:
            return 0.0
